package io.github.wanggit.antrpc.commons.rate;

import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RateLimiterConfig;
import io.github.wanggit.antrpc.commons.bean.RegisterBeanMethod;
import io.vavr.CheckedFunction0;
import io.vavr.control.Try;

import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 *
 *
 * <pre>
 *     https://resilience4j.readme.io/docs/ratelimiter
 * </pre>
 */
public class Resilience4jRateLimiterContainer implements IResilience4jRateLimiter {

    private final ConcurrentHashMap<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, LimitInfo> limitInfos = new ConcurrentHashMap<>();

    @Override
    public Object run(
            RegisterBeanMethod registerBeanMethod,
            CheckedFunction0<Object> nextFunction,
            Consumer<? super Throwable> consumer,
            Function<? super Throwable, ?> onFailFunction) {
        String callUniqueKey = createRateLimiterKey(registerBeanMethod);
        RateLimiter rateLimiter = rateLimiters.get(callUniqueKey);
        if (null != rateLimiter) {
            CheckedFunction0<Object> objectCheckedFunction0 =
                    RateLimiter.decorateCheckedSupplier(rateLimiter, nextFunction);
            return Try.of(objectCheckedFunction0)
                    .recover(onFailFunction)
                    .onFailure(consumer)
                    .getOrElse(() -> null);
        }
        return Try.of(nextFunction)
                .recover(onFailFunction)
                .onFailure(consumer)
                .getOrElse(() -> null);
    }

    @Override
    public LimitInfo getLimitInfo(RegisterBeanMethod registerBeanMethod) {
        return limitInfos.get(createRateLimiterKey(registerBeanMethod));
    }

    @Override
    public void checkOrChangeEventCountCircuitBreaker(RegisterBeanMethod registerBeanMethod) {
        if (registerBeanMethod.getLimit() <= 0 || registerBeanMethod.getDurationInSeconds() <= 0) {
            return;
        }
        String callUniqueKey = createRateLimiterKey(registerBeanMethod);
        if (!rateLimiters.containsKey(callUniqueKey)
                || !LimitInfo.checkEquals(limitInfos.get(callUniqueKey), registerBeanMethod)) {
            synchronized (callUniqueKey.intern()) {
                if (!rateLimiters.containsKey(callUniqueKey)
                        || !LimitInfo.checkEquals(
                                limitInfos.get(callUniqueKey), registerBeanMethod)) {
                    RateLimiterConfig rateLimiterConfig =
                            RateLimiterConfig.custom()
                                    .timeoutDuration(Duration.ofMillis(0))
                                    .limitRefreshPeriod(
                                            Duration.ofSeconds(
                                                    registerBeanMethod.getDurationInSeconds()))
                                    .limitForPeriod(registerBeanMethod.getLimit())
                                    .build();
                    rateLimiters.remove(callUniqueKey);
                    rateLimiters.putIfAbsent(
                            callUniqueKey, RateLimiter.of(callUniqueKey, rateLimiterConfig));
                    limitInfos.remove(callUniqueKey);
                    limitInfos.putIfAbsent(
                            callUniqueKey,
                            new LimitInfo(
                                    registerBeanMethod.getLimit(),
                                    registerBeanMethod.getDurationInSeconds()));
                }
            }
        }
    }

    private String createRateLimiterKey(RegisterBeanMethod registerBeanMethod) {
        return registerBeanMethod.getClassName() + "#" + registerBeanMethod.toString();
    }
}
